"""
向量存储管理模块
"""
import chromadb
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.storage.storage_context import StorageContext
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings
from llama_index.core.response_synthesizers import get_response_synthesizer
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.llms.openai import OpenAI
from document_processor import DocumentProcessor
import uuid

from config import CHROMA_DB_DIR, COLLECTION_NAME

class VectorStoreManager:
    def __init__(self):
        """初始化向量存储管理器"""
        self.reset_database()
        print(f"初始化 ChromaDB 客户端，路径: {CHROMA_DB_DIR}")
        self.chroma_client = chromadb.PersistentClient(path=CHROMA_DB_DIR)
        print(f"获取或创建集合: {COLLECTION_NAME}")
        self.chroma_collection = self.chroma_client.get_or_create_collection(
            name=COLLECTION_NAME,
        )
        self.embed_model = OpenAIEmbedding()
        self.llm = OpenAI(model="gpt-4o")
        Settings.embed_model = self.embed_model
        Settings.llm = self.llm
        # 初始化向量存储
        print("初始化向量存储...")
        self.vector_store = SimpleVectorStore(chroma_collection=self.chroma_collection)
        self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
        print(f"当前集合中的文档数: {self.get_collection_stats()}")

    def clear_collection(self):
        """清空当前集合"""
        try:
            # 删除现有集合
            self.chroma_client.delete_collection(COLLECTION_NAME)
            print(f"已删除集合: {COLLECTION_NAME}")
            
            # 重新创建集合
            self.chroma_collection = self.chroma_client.create_collection(
                name=COLLECTION_NAME,
                metadata={"hnsw:space": "cosine"}
            )
            print(f"已重新创建空集合: {COLLECTION_NAME}")
            
            # 重新初始化向量存储
            self.vector_store = SimpleVectorStore(chroma_collection=self.chroma_collection)
            self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
            
            return True
        except Exception as e:
            print(f"清空集合时出错: {str(e)}")
            return False
            
    def reset_database(self):
        """完全重置数据库（删除所有集合）"""
        try:
            # 获取所有集合
            collections = self.chroma_client.list_collections()
            
            # 删除所有集合
            for collection in collections:
                self.chroma_client.delete_collection(collection.name)
                print(f"已删除集合: {collection.name}")
            
            # 重新创建我们的集合
            self.chroma_collection = self.chroma_client.create_collection(
                name=COLLECTION_NAME,
            )
            print(f"已重新创建空集合: {COLLECTION_NAME}")
            
            # 重新初始化向量存储
            self.vector_store = SimpleVectorStore(chroma_collection=self.chroma_collection)
            self.storage_context = StorageContext.from_defaults(vector_store=self.vector_store)
            
            return True
        except Exception as e:
            print(f"重置数据库时出错: {str(e)}")
            return False

    def create_index(self, files=None):
        """
        创建文档索引
        
        Args:
            files (list, optional): 文件列表
            
        Returns:
            VectorStoreIndex: 创建的索引对象
        """
        document_processor = DocumentProcessor()
        documents = document_processor.process_documents()
        if files:
            print(f"处理上传的文件: {len(files)}个")
            for file in files:
                documents.append(Document(text=file.getvalue().decode(), file_name=file.name))
        
        if documents:
            print(f"开始处理文档，总数: {len(documents)}")
            # 文档统计
            doc_stats = document_processor.get_document_stats()
            print(f"总文件数: {doc_stats['total_files']}")
            print(f"有效文件数: {doc_stats['valid_files']}")
            
            # 显示文件类型统计
            if doc_stats['file_types']:
                print("文件类型分布:")
                for ext, count in doc_stats['file_types'].items():
                    print(f"- {ext}: {count}个文件")
            
            # 向量存储统计
            print("创建索引前的向量存储统计...")
            collection_stats = self.get_collection_stats()
            print(f"当前集合中的文档数: {collection_stats}")

            print("开始创建索引...")
            
            # 直接使用 ChromaDB API 添加文档
            try:
                print("使用 ChromaDB API 直接添加文档...")
                embeddings = []
                texts = []
                ids = []
                metadatas = []
                
                for doc in documents:
                    # 生成文档的嵌入向量
                    embedding = self.embed_model.get_text_embedding(
                        doc.text
                    )
                    
                    # 准备数据
                    embeddings.append(embedding)
                    texts.append(doc.text)
                    ids.append(str(uuid.uuid4()))
                    metadatas.append({"source": doc.metadata.get("file_name", "unknown")})
                
                # 批量添加到 ChromaDB
                self.chroma_collection.add(
                    embeddings=embeddings,
                    documents=texts,
                    ids=ids,
                    metadatas=metadatas
                )
                print(f"成功添加 {len(documents)} 个文档到 ChromaDB")
                
            except Exception as e:
                print(f"添加文档到 ChromaDB 时出错: {str(e)}")
            
            # 创建索引
            index = VectorStoreIndex.from_documents(documents)
            
            print("索引创建完成")
            print(f"创建索引后的向量存储统计: {self.get_collection_stats()}")
            
            return index
        return None
    
    def get_collection_stats(self):
        """获取向量存储统计信息"""
        try:
            count = self.chroma_collection.count()
            print(f"ChromaDB collection count: {count}")
            return count
        except Exception as e:
            print(f"获取集合统计信息时出错: {str(e)}")
            return 0
    
    def query(self, index, query_text, streaming=False):
        """
        查询索引
        
        Args:
            index (VectorStoreIndex): 索引对象
            query_text (str): 查询文本
            streaming (bool): 是否使用流式响应
            
        Returns:
            Response: 查询响应对象
        """

        query_engine = index.as_query_engine()
        
        return query_engine.query(query_text + "，请用中文回答")